const { TFEOpAttr } = require("@tensorflow/tfjs-node/dist/tfjs_binding");
const getData = require("./data");
const tf = require("@tensorflow/tfjs-node");

const TRAIN_DIR = "垃圾分类/train";
const OUTPUT_DIR = "output";
const MOBILENT_URL = "http://ai-sample.oss-cn-hangzhou.aliyuncs.com/pipcook/models/mobilenet/web_model/model.json";

const main = async () => {
    //加载数据
    const { ds, classes } = await getData(TRAIN_DIR, OUTPUT_DIR);
    //定义模型
    const mobilenet = await tf.loadLayersModel(MOBILENT_URL);
    mobilenet.summary();
    //生名一个模型
    const model = tf.sequential();
    //截断神经网络
    for (let i = 0; i < 86; i++) {
        const layer = mobilenet.layers[i];
        layer.trainable = false;
        model.add(layer);
    }
    model.add(tf.layers.flatten());
    //隐藏层
    model.add(tf.layers.dense({
        units:10,
        activation:"relu"//激活函数
    }))
    //输出层
    model.add(tf.layers.dense({
        units:classes.length,
        activation:"softmax"//激活函数
    }))
    //训练模型
    model.compile({
        loss:"sparseCategoricalCrossentropy",//损失函数
        optimizer:tf.train.adam(),//优化器
        metrics:["acc"],//显示当前的准确度
    })
    await model.fitDataset(ds,{
        epochs:20
    });
    await model.save(`file://${process.cwd()}/${OUTPUT_DIR}`);


}

main();